import numpy as np
from easy_io import read_pkl_file, write_pkl_file


def prepare_labelinfo(label_dict):
    noduleids = {}
    centers = {}
    radius = {}
    maligns = {}
    noduleid = 0
    for scanid in sorted(label_dict.keys()):
        nodules = label_dict[scanid]
        if len(nodules) > 0:
            noduleids[scanid] = np.arange(noduleid, noduleid + len(nodules), dtype='int')
            noduleid += len(nodules)
            bbox = np.asarray([n['bbox'] for n in nodules], 'float')
            assert bbox.ndim == 2 and bbox.shape[-1] == 6
            centers[scanid] = (bbox[:, :3] + bbox[:, 3:] - 1) / 2
            radius[scanid] = (bbox[:, 3:] - bbox[:, :3] - 1) / 2 + 2.0
            maligns[scanid] = np.asarray([n['malign'] if 'malign' in n else n['attrs']['malignancy'] for n in nodules], 'float')
    return noduleids, centers, radius, maligns


def prepare_labeler(labelinfo):
    noduleids, centers, radius, maligns = labelinfo

    def labeler(c):
        scanid = c['scanid']
        bbox = np.asarray(c['bbox'], 'float')
        center = (bbox[:3] + bbox[3:] - 1) / 2
        if scanid in noduleids:
            diff = np.sum(((center - centers[scanid]) / radius[scanid]) ** 2, axis=-1) <= 1.0
            return noduleids[scanid][diff], maligns[scanid][diff]
        else:
            return np.empty((0,), 'int'), np.empty((0,), 'float')

    return labeler


def label(candidates, pos_labelinfo, mid_labelinfo=None):
    pos_labeler = prepare_labeler(pos_labelinfo)
    if mid_labelinfo:
        mid_labeler = prepare_labeler(mid_labelinfo)
    for c in candidates:
        noduleidlst, malignlst = pos_labeler(c)
        if len(noduleidlst) > 0:
            if len(noduleidlst) > 1:
                print('Warning! More than one nodule has been found by', c)
            c.update(polarity=1, noduleid=noduleidlst[0], malign=malignlst[0])
        elif mid_labelinfo:
            noduleidlst, malignlst = mid_labeler(c)
            if len(noduleidlst) > 0:
                c.update(polarity=0, noduleid=noduleidlst[0], malign=malignlst[0])
            else:
                c.update(polarity=-1, noduleid=-1, malign=-1.0)
        else:
            c.update(polarity=-1, noduleid=-1, malign=-1.0)
    return candidates


def main(infile, outfile,
         pos_labelinfo_file=None, pos_label_dict_file=None,
         mid_labelinfo_file=None, mid_label_dict_file=None):
    assert pos_labelinfo_file or pos_label_dict_file
    # pos_labelinfo = (pos_labelinfo_file and read_pkl_file(pos_labelinfo_file)) or \
    #                 (pos_label_dict_file and prepare_labelinfo(np.load(pos_label_dict_file).item()))
    # mid_labelinfo = (mid_labelinfo_file and read_pkl_file(mid_labelinfo_file)) or \
    #                 (mid_label_dict_file and prepare_labelinfo(np.load(mid_label_dict_file).item())) or None
    pos_labelinfo = (pos_labelinfo_file and read_pkl_file(pos_labelinfo_file)) or \
                    (pos_label_dict_file and prepare_labelinfo(read_pkl_file(pos_label_dict_file))) or None
    mid_labelinfo = (mid_labelinfo_file and read_pkl_file(mid_labelinfo_file)) or \
                    (mid_label_dict_file and prepare_labelinfo(read_pkl_file(mid_label_dict_file))) or None
    candidates = read_pkl_file(infile)
    candidates = label(candidates, pos_labelinfo, mid_labelinfo)
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v3.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v4.pkl',
    #     pos_label_dict_file='/data_4t/Kaggle/lidc&kaggle/label_dict3.npy',
    #     mid_label_dict_file='/data_4t/Kaggle/lidc&kaggle/label_dict1.npy'
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v2.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v2_revised.pkl',
    #     pos_label_dict_file='/data_4t/Kaggle/lidc&kaggle/label_dict3.npy',
    #     mid_label_dict_file='/data_4t/Kaggle/lidc&kaggle/label_dict1.npy'
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v8.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v8_revised.pkl',
    #     pos_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_3_backup.pkl',
    #     mid_label_dict_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_label_dict_1_backup.pkl',
    # )
    # lidc_3 = np.load('/data_4t/Kaggle/backup/lidc/label_dict_3_0329.npy').item()
    # kaggle = np.load('/data_4t/Kaggle/backup/kaggle/label_dict_0325.npy').item()
    # label_dict_3 = dict()
    # label_dict_3.update(lidc_3)
    # label_dict_3.update(kaggle)
    # write_pkl_file('/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #                prepare_labelinfo(label_dict_3))
    # lidc_1 = np.load('/data_4t/Kaggle/backup/lidc/label_dict_1_0328.npy').item()
    # label_dict_1 = lidc_1
    # write_pkl_file('/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    #                prepare_labelinfo(label_dict_1))
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
    # main(
    #     infile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v2.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v2_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_kaggle_testset_candidates_v4.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_kaggle_testset_candidates_v4_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )

    main(
        infile='/data_4t/Kaggle/candidates/3dcnn_spie_candidates_v4.pkl',
        outfile='/data_4t/Kaggle/candidates/3dcnn_spie_candidates_v4_newlabel.pkl',
        pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl'
    )

    # write_pkl_file('/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl',
    #                prepare_labelinfo(np.load('/data_4t/Kaggle/backup/spie/label_dict.npy').item()))

    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble_kaggle_testset_candidates_v4.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble_kaggle_testset_candidates_v4_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )

    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble_spie_candidates_v4.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_ensemble_spie_candidates_v4_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/spie_labelinfo.pkl',
    # )

    # main(
    #     infile='/data_4t/Kaggle/candidates/3dcnn_lidc_kaggle_candidates_v9_nomerge_nms.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_lidc_kaggle_candidates_v9_nomerge_nms_newlabel.pkl',
    #     pos_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_3.pkl',
    #     mid_labelinfo_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_labelinfo_1.pkl',
    # )
